from model import *
from prompt import *
from fastapi import FastAPI, Header, BackgroundTasks, HTTPException
from fastapi.responses import FileResponse
from fine_tune_service import *
import uvicorn

app = FastAPI()
app.config = { 'fine_tune_status': 'none' }

@app.post('/gptbots/{botId}/clarifyprompt')
async def get_clarify_prompt(botId: str, input: ClarifyPromptInput) -> ClarifyPromptOutput:
    prompts = get_prompt()
    return ClarifyPromptOutput(type='noAnswer', prompt=prompts)

@app.post('/gptbots/{botId}/queryprompt')
async def get_query_prompt(input: QueryPromptInput, botId: str) -> QueryPromptOutput:
    prompts = get_prompt()
    final_srcs = get_final_srcs()
    print([{ 'number': index +1, 'title': src['title'], 'url': src['url']} for index, src in enumerate(final_srcs)])
    return QueryPromptOutput(prompt=prompts, articles=[{ 'number': index +1, 'title': src['title'], 'url': src['url']} for index, src in enumerate(final_srcs)])

@app.post('/gptbots/{botId}/query')
async def query(botId: str, input: QueryInput) -> QueryOutput:
    print([{ 'role': p.role, 'content': p.content } for p in input.prompt])
    return QueryOutput(prompt=input.prompt, articles=input.articles)

@app.get('/fine_tune')
async def fine_tune(background_task: BackgroundTasks):
    fine_tune_status = app.config['fine_tune_status']
    if fine_tune_status == 'pending' or fine_tune_status == 'running':
        return { 'status': fine_tune_status, 'msg': '' }
    background_task.add_task(fine_tune_task)
    return { 'status': 'pending', 'msg': '' }

@app.get('/fine_tune_result/download')
async def fine_tune_download():

     # Specify the path to the file you want to download
    file_path = 'fine_tune_result.json'
    
    # Specify the filename for the client to download
    filename = 'fine_tune_result.json'
    
    # Return the file to the client for download
    return FileResponse(file_path, filename=filename)

@app.get('/fine_tune/status')
async def fine_tune_status() -> None:
    fine_tune_status = app.config['fine_tune_status']
    return { 'status': fine_tune_status, 'msg': '' }

@app.get('/articles')
async def get_aritics() -> List[Article]:
    articles = [
        {
            "number": 1,
            "title": "1",
            "url": "url1",
            "score": 0.82
        },
        {
            "number": 1,
            "title": "1",
            "url": "url1",
            "score": 0.81
        },
        {
            "number": 1,
            "title": "1",
            "url": "url1",
            "score": 0.83
        },
        {
            "number": 1,
            "title": "1",
            "url": "url1",
            "score": 0.86
        }
    ]
    articles.sort(key=get_score)
    articles.reverse()
    print('before' , articles)
    topArticles = 3
    if topArticles > 0 and len(articles) > topArticles:
        articles = articles[:topArticles]
        print('after', articles)

    return articles

def get_score(e):
    return e['score']

def fine_tune_task():
    print('task start')

    app.config['fine_tune_status'] = 'running'

    status = 'running'
    msg = ''

    try:
        exec_fine_tune(app)
        
    except Exception as ex:
        print(ex)
        status = 'error'
        msg = ex.args
    
    return { 'status': status, 'msg': msg }

def get_final_srcs():
    return [{
        "file": "file1",
        "title": "title1",
        "url": "url1",
        "segment": "segment1",
        "score": 0.2
    },
    {
        "file": "file2",
        "title": "title3",
        "url": "url3",
        "segment": "segment3",
        "score": 0.3
    }]

if __name__ == '__main__':
    uvicorn.run(app, host="127.0.0.1", port=8000)